#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/15
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
import torch
from four_four_build_network import MYNET
import torchvision.transforms as transforms
from PIL import Image

# 泛化测试
# 测试图像处理
transform = transforms.Compose([transforms.Resize((64, 64)),
                                transforms.Grayscale(),
                                transforms.ToTensor()])

# 加载模型
model = MYNET()
# 注意：torch.load('tmp/4-model.pkl')返回的是一个状态字典，而不是模型本身。
# 状态字典包含了模型的参数和训练状态，可以用来恢复模型的状态。
model.load_state_dict(torch.load('tmp/4-model.pkl'))
model.eval()

# 输入图像并预测
img = Image.open('data4/test/00007/3544.png')
img = transform(img)
img = img.view(1, 1, 64, 64)
output = model(img)
# 使用torch.max函数获取output中每一行最大值的索引，返回两个值，第一个是最大值，第二个是最大值的索引，这里我们只关心最大值的索引
_, prediction = torch.max(output, 1)
# 将最大值的索引从tensor转换为numpy数组
prediction = prediction.numpy()[0]
print(prediction)
